import cv2
import os

input_path = './val2014'
txt_path = './detection-results'
output_path= './rectangle'

IMG_EXT = ['.jpg', '.JPG', '.png', '.PNG', '.bmp', '.BMP', '.jpeg', '.JPEG']
images_list = [os.path.join(input_path, img) for img in os.listdir(input_path) if os.path.splitext(img)[1] in IMG_EXT]

if not os.path.exists(output_path):
    os.makedirs(output_path)

for image_path in images_list:
    image = cv2.imread(image_path)

    for line in open(os.path.join(txt_path, os.path.splitext(os.path.basename(image_path))[0] + '.txt')): 
        category = line.split(" ")[0]
        score = line.split(" ")[1]
        x1 = int(eval(line.split(" ")[2]))
        y1 = int(eval(line.split(" ")[3]))
        x2 = int(eval(line.split(" ")[4]))
        y2 = int(eval(line.split(" ")[5]))
        first_point = (x1, y1)
        last_point = (x2, y2)
        cv2.rectangle(image, first_point, last_point, (0, 255, 0), 2)
        cv2.putText(image, category + "--" + score[0:8], (x1+100,y1+100), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,255), 2)
    
    cv2.imwrite(os.path.join(output_path, os.path.basename(image_path)), image)
